#ifndef __GRU_H__
#define __GRU_H__

#include "Linear.h"
#include "ProcessingUnit.h"

typedef struct {
    float* w;
    float b;
} LinearInitStruct_t;

typedef struct {
    float* w_ir;
    float* w_iz;
    float* w_in;
    float* w_hr;
    float* w_hz;
    float* w_hn;
    float* b_ir;
    float* b_iz;
    float* b_in;
    float* b_hr;
    float* b_hz;
    float* b_hn;
} GruInitStruct_t;

class GRUCell {
public:
    GRUCell() {};
    GRUCell(const int hidden_size);
    void process(const Eigen::VectorXf& input, Eigen::VectorXf& output);
    Eigen::VectorXf getHiddenState() const { return _hidden; };
    void setParams(float* w_ir, float* w_iz, float* w_in, float* w_hr,
        float* w_hz, float* w_hn, float* b_ir, float* b_iz,
        float* b_in, float* b_hr, float* b_hz, float* b_hn);
    void reset();

protected:
    int _input_size, _hidden_size;
    Eigen::MatrixXf _hidden;
    Eigen::MatrixXf _w_ir;
    Eigen::MatrixXf _w_iz;
    Eigen::MatrixXf _w_in;
    Eigen::MatrixXf _w_hr;
    Eigen::MatrixXf _w_hz;
    Eigen::MatrixXf _w_hn;
    Eigen::VectorXf _b_ir;
    Eigen::VectorXf _b_iz;
    Eigen::VectorXf _b_in;
    Eigen::VectorXf _b_hr;
    Eigen::VectorXf _b_hz;
    Eigen::VectorXf _b_hn;
    Eigen::VectorXf _r;
    Eigen::VectorXf _z;
    Eigen::VectorXf _n;
};

class GRU : public ProcessingUnit {
public:
    GRU() {};
    GRU(const int hidden_size);
    void process(const float* input, float* output, size_t num_samples) override;
    void setParams(const std::string name, void* params);
    void reset();

protected:
    int _hidden_size;
    Eigen::VectorXf _head_weight;
    float _head_bias;
    GRUCell _rec;
    Linear _lin;
    Eigen::VectorXf _input, _hidden, _output;

    float _processSample(const float x);
};

#endif